#include <os/task.h>
#include <os/pthread.h>
#include <os/process.h>
#include <os/schedule.h>
#include <os/debug.h>
#include <os/spinlock.h>
#include <arch/mp.h>
#include <lib/assert.h>
#include <lib/stdlib.h>

#ifdef ENABLE_SMP

// dispatch a task to suiface schedule unit
void TaskDispatch(task_t *task)
{
    // record current cpu id
    task->last_cpuid = CpuGetMyId();

    processor_info_t *proc = MpGetProcessor();
    if (!proc)
    {
        Panic("[dispatch] get processor failed!\n");
    }

    sched_unit_t *su = SchedGetCurUnit();
    if (!su)
    {
        Panic("get current schedule unit failed!\n");
    }

    if (TASK_IS_KERNEL_THREAD(task)) // task is thread
    {
        if (task->last_cpuid >= 0) // thread dispatch to last exec processor to improve performance with cache
        {
            proc = MpGetProcessorById(task->last_cpuid);
            if (!proc)
            {
                task_t *task = TaskFindByPid(task->processgroup_id); // find processs task
                proc = MpGetProcessorById(task->cpuid);
                if (!proc)
                {
                    su = SchedGetFreeUnit(); // prcess current processor is undef,try get free su unit
                }
                su = SchedFindByCpu(task->cpuid);
            }
            su = SchedFindByCpu(task->last_cpuid);
        }
        else
        {
            su = SchedGetFreeUnit(); // swicth a free unit
        }
    }
    else
    {
        // task is process
        if (task->last_cpuid >= 0)
        {
            proc = MpGetProcessorById(task->last_cpuid);
            if (!proc)
            {
                su = SchedGetFreeUnit(); // switch a free unit
            }
            su = SchedFindByCpu(task->last_cpuid); // dispatch last unit
        }
        else
        {
            su = SchedGetFreeUnit(); // switch a free unit
        }
    }

    // dispatch to current unit
    if (su->tasknum < TASK_NUM_HIGH)
    {
        SchedQueueAddTaskTail(su, task);
        return;
    }

    // dispatch to free unit
    uint8_t prior = task->static_priority;
    if (prior != TASK_PRIOR_LEVEL_HIGH)
    {
        sched_unit_t *free_su = SchedGetFreeUnit(); // get free unit
        if (!free_su)
        {
            SchedQueueAddTaskTail(su, task);
        }
        // dispatch to free unit
        SchedQueueAddTaskTail(free_su, task);
    }
}

void TaskMove(task_t *task, cpuid_t cpu)
{
    // no move high or realtime task
    if (task->priority != TASK_PRIOR_LEVEL_HIGH || task->priority != TASK_PRIOR_LEVEL_REALTIME)
        return;

    // task in blocked status direct add to new schedule unit
    if (task->status == TASK_BLOCKED)
    {
        SchedQueueAddTaskTail(SchedFindByCpu(cpu), task);
    }

    // remove from old schedule unit
    SchedQueueRemoveTask(SchedGetCurUnit(), task);
    // add to new schedule unit
    SchedQueueAddTaskTail(SchedFindByCpu(cpu), task);
}

// called by current schedule unit to push task to new su
void TaskBalancePush()
{
    int i;
    task_t *task = NULL;
    float total = schedule->tasknum;
    sched_unit_t *su = SchedGetCurUnit();
    if (!su)
    {
        Panic("Schedule error!\n");
    }
    float tasks = su->tasknum;
    float load = tasks / total;

    // check cpu load to switch balance method
    // in current,we olny move that tasks in running on loads >= CPU_TASK_LOAD_MAX
    // and we don't move that no ready tasks
    if (load >= CPU_TASK_LOAD_MAX)
    {
        sched_unit_t *free = SchedGetFreeUnit();
        if (!free)
            return;

        list_traversal_all_owner_to_next(task, &task_global_list, list)
        {
            if (task->cpuid != su->cpuid) // no itselt unit tasks
                continue;

            if (task->status >= TASK_BLOCKED) // no ready tasks
                continue;

            if (task->status != TASK_RUNNING) // task no running
            {
                TaskMove(task, free->cpuid); // move to free cpu
            }
            else
            {
                if (task->status == TASK_RUNNING) // task in running
                {
                    TaskMove(task, free->cpuid);
                    TaskReady(task); // set task to ready
                    Schedule();      // schedule next task
                }
            }
        }
    }
    else
    {
        if (load >= CPU_TASK_LOAD_WARN)
        {
            sched_unit_t *free = SchedGetFreeUnit();

            list_traversal_all_owner_to_next(task, &task_global_list, list)
            {
                if (task->cpuid != su->cpuid) // no itselt unit tasks
                    continue;

                if (task->priority >= TASK_PRIOR_LEVEL_HIGH) // don't to move realtime and high priority tasks
                    continue;

                if (task->status == TASK_RUNNING) // don't to move in running tasks
                    continue;

                if (task->status >= TASK_BLOCKED) // don'tv to move no ready tasks
                    continue;

                sched_unit_t *free = SchedGetFreeUnit();
                if (!free)
                    free = su;
                TaskMove(task, free->cpuid);
            }
        }
    }
}

// called by current unit that when current schedule unit free to pull a task
void TaskBalancePull()
{
    sched_unit_t *su = SchedGetCurUnit();
    if (!su)
    {
        Panic("schedule unit error!\n");
    }
    float load = su->tasknum / schedule->tasknum;
    // curent cpu load too high
    if (load >= CPU_TASK_LOAD_WARN)
        return;

    int i;
    sched_unit_t *unit = schedule->sched_unit_table;
    sched_unit_t *busy = NULL;
    float t_load;

    // find busy unit
    for (i = 0; i < CPU_NUM_MAX; i++)
    {
        unit = &schedule->sched_unit_table[i];
        // skip itselt unit
        if (unit->cpuid == su->cpuid)
            continue;
        t_load = su->tasknum / unit->tasknum;
        if (t_load > load) // record busy unit
        {
            busy = unit;
        }
    }

    // current unit no busy
    if (busy != su)
    {
        task_t *task;
        sched_queue_t *queue = &busy->priority_queue[TASK_PRIOR_LEVEL_NORMAL];
        list_traversal_all_owner_to_next(task, &task_global_list, list)
        {
            // don't move running task and not ready task
            if (task->status == TASK_RUNNING)
                continue;

            if (task->status >= TASK_BLOCKED)
                continue;

            // task move to current unit
            TaskMove(task, su->cpuid);
        }
    }
}

static void TaskBalanceHandler(void *arg)
{
    // push and pop task
    TaskBalancePush();
    TaskBalancePull();

    timer_t *timer = (timer_t *)arg;
    timer->timeout = TASK_BALANCE_TIMEOUT;
    TimerAdd(timer);
}

void TaskBalanceInit()
{
    // register load balance timer
    timer_t *timer = TimerAlloc();
    if (!timer)
    {
        KMemAlloc("no mem for balance timer\n");
        return;
    }
    TimerInit(timer, TASK_BALANCE_TIMEOUT, timer, TaskBalanceHandler);
}

#endif